﻿#include "SimpleShapeBasedMatching.h"
//#include <QDir>
//#include <QImage>
#include <QtDebug>

namespace Qly {

SimpleShapeBasedMatching::SimpleShapeBasedMatching()
{
    shape_based_matching::shapeInfo_producer::Info info(0, 1, 0);
    infos_have_templ.push_back(info);
}

bool SimpleShapeBasedMatching::trainModel(cv::Mat modelImg, cv::Mat mask, int num_features)
{
    if(modelImg.empty() || mask.size() != modelImg.size() || mask.channels() != 1 || num_features < 3)
    {
        return false;
    }
    if(modelImg.channels() == 4)
    {
        cv::cvtColor(modelImg, templateImg, cv::COLOR_BGRA2BGR);
    }
    else
    {
        templateImg = modelImg.clone();
    }
    padded_templateImg = templateImg.clone();
    padded_mask = mask.clone();
    if(pDetector)
    {
        delete pDetector;
    }
    pDetector = new line2Dup::Detector(num_features, {4, 8});   
    infos_have_templ[0].templ_id = pDetector->addTemplate(padded_templateImg, class_id, padded_mask, num_features);
    if(infos_have_templ[0].templ_id == -1)
    {
        return false;
    }

    return true;
}

bool SimpleShapeBasedMatching::trainModel(cv::Mat modelImg, int num_features)
{
    preprocessModelImage(modelImg, 0, 0);
    cv::Mat mask = cv::Mat(templateImg.size(), CV_8UC1, {255});
    if(pDetector)
    {
        delete pDetector;
    }
    pDetector = new line2Dup::Detector(num_features, {4, 8});
    infos_have_templ[0].templ_id = pDetector->addTemplate(padded_templateImg, class_id, padded_mask, num_features);

    if(infos_have_templ[0].templ_id == -1)
    {
        return false;
    }

    return true;
}

bool SimpleShapeBasedMatching::saveModel(std::string path, std::string modelName)
{
    if(pDetector == nullptr) return false;
    std::string fullModelName = path + "/" + modelName + ".yaml";
    std::string fullModelImageName = path + "/" + modelName + ".png";
    // save templates
    pDetector->writeClasses(fullModelName);
    return cv::imwrite(fullModelImageName, templateImg);
}

bool SimpleShapeBasedMatching::loadModel(std::string path, std::string modelName)
{
    std::string fullModelName = path + "/" + modelName + ".yaml";
    std::string fullModelImageName = path + "/" + modelName + ".png";
    if(pDetector)
    {
        delete pDetector;
    }
    pDetector = new line2Dup::Detector(64, {4, 8});
    templateImg = cv::imread(fullModelImageName);
    pDetector->readClasses(class_ids, fullModelName);

    return !templateImg.empty();
}

std::vector<ShapeBasedMatching::MatchResult> SimpleShapeBasedMatching::runMatch(const cv::Mat &img, float score, size_t maxCount)
{
    std::vector<ShapeBasedMatching::MatchResult> ret;
    if(img.empty() || !pDetector) return ret;
    resultImg = img.clone();
    cv::Mat mat = preprocessTestImage(img);
    std::vector<line2Dup::Match> matches = pDetector->match(mat, score, class_ids);// 这里得到的 matches 已经是排好序的了

    qDebug() <<"pDetector->match size = " << matches.size();
    size_t count = std::min(maxCount, matches.size());
    for(size_t i = 0; i < count; i ++)
    {
        MatchResult mr;
        const line2Dup::Match &match = matches[i];
        mr.rect = rotatedRect(match);
        mr.score = match.similarity;
        ret.push_back(mr);
    }
    return ret;
}

}
